#pragma once
#include "ShaderMap.h"
#include "RHIShader.h"
#include "../../Core/CoreGlobals.h"

template<typename ShaderType>
class ShaderPtr
{

public:

	void ActiveShader(RHIDescriptor* descriptor)
	{
		assert(Reference->mDescriptor == nullptr);
		Reference->SetDescriptor(descriptor);
	}

	__forceinline ShaderPtr()
		: Reference(nullptr)
	{
	}

	__forceinline ShaderPtr(uint32 mask)
	{
		Reference = NewObject<ShaderAdapter>();
		Reference->SetShader(GShaderMap->GetShader<ShaderType>());
		if (Reference)
		{
			Reference->AddRef();
		}
	}

	__forceinline ShaderPtr(const std::string& name)
	{
		Reference = NewObject<ShaderAdapter>();
		Reference->SetShader(GShaderMap->GetShader(name));
		if (Reference)
		{
			Reference->AddRef();
		}
	}


	ShaderPtr(const ShaderPtr& Copy)
	{
		Reference = Copy.Reference;
		if (Reference)
		{
			Reference->AddRef();
		}
	}

	//template<typename CopyShaderType>
	//explicit ShaderPtr(const ShaderPtr<CopyShaderType>& Copy)
	//{
	//	assert(Reference == nullptr);
	//	Reference = static_cast<ShaderAdapter*>(Copy.GetReference());
	//	if (Reference)
	//	{
	//		Reference->AddRef();
	//	}
	//}

	template<typename CopyShaderType>
	__forceinline ShaderPtr(const ShaderPtr<CopyShaderType>& Copy)
	{
		Reference = static_cast<ShaderAdapter*>(Copy.GetReference());
		if (Reference)
		{
			Reference->AddRef();
		}
	}

	__forceinline ShaderPtr(ShaderPtr&& Move)
	{
		Reference = Move.Reference;
		Move.Reference = nullptr;
	}

	template<typename MoveShaderType>
	explicit ShaderPtr(ShaderPtr<MoveShaderType>&& Move)
	{
		Reference = Move.GetReference();
		Move.Reference = nullptr;
	}

	~ShaderPtr()
	{
		if (Reference)
		{
			if (Reference->Release() == 0)
				DeleteObject<ShaderAdapter>(Reference);
		}
	}

	ShaderPtr& operator=(ShaderAdapter* InReference)
	{
		// Call AddRef before Release, in case the new reference is the same as the old reference.
		ShaderAdapter* OldReference = Reference;
		Reference = InReference;
		if (Reference)
		{
			Reference->AddRef();
		}
		if (OldReference)
		{
			if (OldReference->Release() == 0)
				DeleteObject<ShaderAdapter>(OldReference);
		}
		return *this;
	}

	__forceinline ShaderPtr& operator=(const ShaderPtr& InPtr)
	{
		return *this = InPtr.Reference;
	}

	template<typename CopyShaderType>
	__forceinline ShaderPtr& operator=(const ShaderPtr<CopyShaderType>& InPtr)
	{
		return *this = InPtr.GetReference();
	}

	ShaderPtr& operator=(ShaderPtr&& InPtr)
	{
		if (this != &InPtr)
		{
			ShaderAdapter* OldReference = Reference;
			Reference = InPtr.Reference;
			InPtr.Reference = nullptr;
			if (OldReference)
			{
				if (OldReference->Release() == 0)
					DeleteObject<ShaderAdapter>(OldReference);
			}
		}
		return *this;
	}

	ShaderType* operator->() const
	{
		assert(Reference != nullptr && Reference->mShader != nullptr);
		Reference->mShader->SetShaderAdapter(Reference);
		ShaderType* shader = static_cast<ShaderType*>(Reference->mShader);
		return shader;
	}

	__forceinline friend uint32 GetTypeHash(const ShaderPtr& InPtr)
	{
		return GetTypeHash(InPtr.Reference);
	}

	__forceinline ShaderAdapter* GetReference() const
	{
		return Reference;
	}

	__forceinline friend bool IsValidRef(const ShaderPtr& InReference)
	{
		return InReference.Reference != nullptr;
	}

	__forceinline bool IsValid() const
	{
		return Reference != nullptr && Reference->mShader != nullptr && Reference->mDescriptor != nullptr;
	}

	uint32 GetRefCount()
	{
		uint32 Result = 0;
		if (Reference)
		{
			Result = Reference->GetRefCount();
			assert(Result > 0); // you should never have a zero ref count if there is a live ref counted pointer (*this is live)
		}
		return Result;
	}

	__forceinline void Swap(ShaderPtr& InPtr) // this does not change the reference count, and so is faster
	{
		ShaderAdapter* OldReference = Reference;
		Reference = InPtr.Reference;
		InPtr.Reference = OldReference;
	}

private:

	ShaderAdapter* Reference;

	template <typename OtherType>
	friend class ShaderPtr;
};

template<typename ShaderType>
__forceinline bool operator==(const ShaderPtr<ShaderType>& A, const ShaderPtr<ShaderType>& B)
{
	return A.GetReference()->mShader == B.GetReference()->mShader && A.GetReference()->mDescriptor == B.GetReference()->mDescriptor;
}

template<typename ShaderType>
__forceinline bool operator==(const ShaderPtr<ShaderType>& A, ShaderType* B)
{
	return A.GetReference()->mShader == B;
}

template<typename ShaderType>
__forceinline bool operator==(ShaderType* A, const ShaderPtr<ShaderType>& B)
{
	return A == B.GetReference()->mShader;
}
